import torch
from model.TextClassifier import Classification
from loguru import logger

if __name__ == '__main__':
    src_len = 7
    batch_szie = 2
    d_model = 32
    num_head = 4
    src = torch.tensor([[4, 3, 2, 6, 0, 0, 0],
                        [5, 7, 8, 2, 4, 0, 0]]).transpose(0, 1)

    src_key_padding_mask = torch.tensor([
        [False, False, False, False, False, True, True],
        [False, False, False, False, False, True, True]])

    model = Classification(vocab_size=10,
                           d_model=d_model,
                           nhead=num_head)
    logits = model(src, src_key_padding_mask=src_key_padding_mask)

    logger.info(logits.shape)
